"""
Companion Recommendation Engine
"""

from sentence_transformers import SentenceTransformer
import numpy as np
from typing import List, Dict, Optional
import time
import os
from collections import defaultdict
from .database import Database


class CompanionRecommender:
    
    def __init__(self, data_dir: str = "data"):
        print("[INFO] Initializing recommender engine...")
        start_time = time.time()
        
        db_path = os.path.join(data_dir, "companion.db")
        self.db = Database(db_path)
        
        self.model = SentenceTransformer('shibing624/text2vec-base-chinese')
        
        self.players = {}
        self.users = {}
        self.orders = []
        
        self.player_embeddings = {}
        
        self.stats = {
            'hot_24h': defaultdict(int),
            'player_status': {},
            'player_clicks': defaultdict(int),
            'player_orders': defaultdict(int),
        }
        
        self.load_data_from_db()
        
        print(f"[INFO] Initialization completed in {time.time() - start_time:.2f}s")
    
    def load_data_from_db(self):
        players = self.db.get_all_players()
        self.load_players(players)
        
        users = self.db.get_all_users()
        self.load_users(users)
        
        orders = self.db.get_all_orders()
        self.load_orders(orders)
    
    def load_players(self, players: List[Dict]):
        for player in players:
            player_id = player['player_id']
            self.players[player_id] = player
            self.stats['player_status'][player_id] = player.get('online', True)
            self.stats['player_orders'][player_id] = player.get('total_orders', 0)
        
        self._build_player_embeddings()
    
    def _build_player_embeddings(self):
        texts = []
        player_ids = []
        
        for player_id, player in self.players.items():
            text = self._build_player_text(player)
            texts.append(text)
            player_ids.append(player_id)
        
        embeddings = self.model.encode(texts, batch_size=32, show_progress_bar=False)
        
        for player_id, emb in zip(player_ids, embeddings):
            self.player_embeddings[player_id] = emb
    
    def _build_player_text(self, player: Dict) -> str:
        parts = [
            player.get('nickname', ''),
            ' '.join(player.get('games', [])),
            player.get('voice_type', ''),
            player.get('description', ''),
            ' '.join(player.get('tags', [])),
            f"价格{player.get('price', 0)}元"
        ]
        return ' '.join(filter(None, parts))
    
    def recommend_home(
        self, 
        user_id: int, 
        page_size: int = 20,
        location: Optional[str] = None
    ) -> List[Dict]:
        is_new_user = self._is_new_user(user_id)
        
        if is_new_user:
            results = self._recommend_for_cold_start(user_id, page_size)
        else:
            results = self._recommend_for_normal_user(user_id, page_size, location)
        
        # Ensure we always return enough results
        if len(results) < page_size:
            results = self._fill_results(results, page_size, user_id)
        
        return results
    
    def _is_new_user(self, user_id: int) -> bool:
        user = self.users.get(user_id, {})
        order_count = user.get('order_count', 0)
        return order_count == 0
    
    def _recommend_for_cold_start(self, user_id: int, page_size: int) -> List[Dict]:
        user = self.users.get(user_id, {})
        interest_games = user.get('interest_games', ['王者荣耀', '和平精英'])
        
        candidates = []
        
        for game in interest_games:
            game_players = self._get_players_by_game(game, top_k=30)
            candidates.extend(game_players)
        
        hot_players = self._get_hot_players(top_k=20)
        candidates.extend(hot_players)
        
        seen = set()
        unique_candidates = []
        for p in candidates:
            if p['player_id'] not in seen:
                seen.add(p['player_id'])
                unique_candidates.append(p)
        
        filtered = self._apply_filters(unique_candidates, user_id)
        filtered.sort(key=lambda x: x['rating'], reverse=True)
        results = self._diversify(filtered, strategy='cold_start')
        
        return results[:page_size]
    
    def _recommend_for_normal_user(
        self, 
        user_id: int, 
        page_size: int,
        location: Optional[str]
    ) -> List[Dict]:
        user_profile = self._get_user_profile(user_id)
        candidates = self._multi_recall(user_id, user_profile)
        filtered = self._apply_filters(candidates, user_id, location)
        ranked = self._rank_candidates(filtered, user_profile)
        diversified = self._diversify(ranked)
        final = self._apply_operation_strategy(diversified, user_profile)
        
        return final[:page_size]
    
    def _get_user_profile(self, user_id: int) -> Dict:
        user = self.users.get(user_id, {})
        user_orders = [o for o in self.orders if o['user_id'] == user_id]
        
        profile = {
            'favorite_games': user.get('favorite_games', ['王者荣耀']),
            'avg_price': user.get('avg_price', 35),
            'preferred_voice': user.get('preferred_voice', '御姐'),
            'order_count': len(user_orders),
            'recent_players': user.get('recent_players', [])
        }
        
        return profile
    
    def _multi_recall(self, user_id: int, user_profile: Dict) -> List[Dict]:
        all_candidates = []
        
        ai_candidates = self._ai_recall(user_id, user_profile, top_k=300)
        all_candidates.extend(ai_candidates)
        
        hot_candidates = self._get_hot_players(top_k=100)
        all_candidates.extend(hot_candidates)
        
        game_candidates = self._game_match_recall(user_profile, top_k=50)
        all_candidates.extend(game_candidates)
        
        seen = set()
        unique = []
        for p in all_candidates:
            if p['player_id'] not in seen:
                seen.add(p['player_id'])
                unique.append(p)
        
        return unique
    
    def _ai_recall(self, user_id: int, user_profile: Dict, top_k: int) -> List[Dict]:
        query_text = self._build_user_query(user_profile)
        query_vec = self.model.encode(query_text)
        
        scores = {}
        for player_id, player_vec in self.player_embeddings.items():
            score = np.dot(query_vec, player_vec) / (
                np.linalg.norm(query_vec) * np.linalg.norm(player_vec)
            )
            scores[player_id] = score
        
        sorted_ids = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:top_k]
        
        candidates = []
        for player_id, score in sorted_ids:
            player = self.players[player_id].copy()
            player['ai_score'] = float(score)
            candidates.append(player)
        
        return candidates
    
    def _build_user_query(self, user_profile: Dict) -> str:
        parts = [
            ' '.join(user_profile.get('favorite_games', [])),
            user_profile.get('preferred_voice', ''),
            f"价格{user_profile.get('avg_price', 30)}元左右"
        ]
        return ' '.join(filter(None, parts))
    
    def _get_hot_players(self, top_k: int = 100) -> List[Dict]:
        sorted_players = sorted(
            self.players.items(),
            key=lambda x: self.stats['player_orders'].get(x[0], 0),
            reverse=True
        )[:top_k]
        
        return [p[1].copy() for p in sorted_players]
    
    def _game_match_recall(self, user_profile: Dict, top_k: int) -> List[Dict]:
        favorite_games = user_profile.get('favorite_games', [])
        if not favorite_games:
            return []
        
        game = favorite_games[0]
        return self._get_players_by_game(game, top_k)
    
    def _get_players_by_game(self, game: str, top_k: int) -> List[Dict]:
        matched = []
        for player in self.players.values():
            if game in player.get('games', []):
                matched.append(player.copy())
        
        matched.sort(key=lambda x: x.get('rating', 0), reverse=True)
        return matched[:top_k]
    
    def _apply_filters(
        self, 
        candidates: List[Dict], 
        user_id: int,
        location: Optional[str] = None
    ) -> List[Dict]:
        filtered = candidates
        
        filtered = [p for p in filtered if self.stats['player_status'].get(p['player_id'], True)]
        
        user = self.users.get(user_id, {})
        blacklist = user.get('blacklist', [])
        filtered = [p for p in filtered if p['player_id'] not in blacklist]
        
        avg_price = user.get('avg_price', 50)
        price_min, price_max = avg_price * 0.7, avg_price * 1.3
        filtered = [p for p in filtered if price_min <= p.get('price', 0) <= price_max]
        
        recent_players = user.get('recent_players', [])
        filtered = [p for p in filtered if p['player_id'] not in recent_players]
        
        filtered = [p for p in filtered if not p.get('is_blocked', False)]
        
        return filtered
    
    def _fill_results(self, results: List[Dict], page_size: int, user_id: int) -> List[Dict]:
        """Fill results to reach page_size if not enough"""
        if len(results) >= page_size:
            return results
        
        existing_ids = {p['player_id'] for p in results}
        user = self.users.get(user_id, {})
        avg_price = user.get('avg_price', 50)
        price_min, price_max = avg_price * 0.5, avg_price * 1.5
        
        # Get additional candidates from all players
        additional = []
        for player in self.players.values():
            if (player['player_id'] not in existing_ids and 
                price_min <= player.get('price', 0) <= price_max):
                additional.append(player.copy())
        
        # Sort by rating and take what we need
        additional.sort(key=lambda x: x.get('rating', 0), reverse=True)
        needed = page_size - len(results)
        results.extend(additional[:needed])
        
        return results
    
    def _rank_candidates(self, candidates: List[Dict], user_profile: Dict) -> List[Dict]:
        for candidate in candidates:
            ai_score = candidate.get('ai_score', 0.5)
            rating_score = candidate.get('rating', 0) / 5.0
            accept_score = candidate.get('accept_rate', 0.8)
            response_score = 1.0 - min(candidate.get('response_time', 5) / 30, 1.0)
            fresh_score = 1.0 if candidate.get('is_new', False) else 0.5
            
            final_score = (
                ai_score * 0.5 +
                rating_score * 0.2 +
                accept_score * 0.15 +
                response_score * 0.1 +
                fresh_score * 0.05
            )
            
            candidate['final_score'] = final_score
        
        candidates.sort(key=lambda x: x['final_score'], reverse=True)
        
        return candidates
    
    def _diversify(self, candidates: List[Dict], strategy: str = 'normal') -> List[Dict]:
        if strategy == 'cold_start':
            return self._diversify_by_game(candidates)
        
        result = []
        game_count = defaultdict(int)
        price_ranges = {'low': 0, 'mid': 0, 'high': 0}
        voice_count = defaultdict(int)
        new_player_count = 0
        
        for candidate in candidates:
            main_game = candidate.get('games', [''])[0]
            if game_count[main_game] >= 5:
                continue
            
            price = candidate.get('price', 0)
            price_range = 'low' if price < 25 else ('high' if price > 45 else 'mid')
            if price_ranges[price_range] >= 8:
                continue
            
            voice = candidate.get('voice_type', '')
            if voice_count[voice] >= 6:
                continue
            
            if candidate.get('is_new') and new_player_count >= 3:
                continue
            
            result.append(candidate)
            game_count[main_game] += 1
            price_ranges[price_range] += 1
            voice_count[voice] += 1
            if candidate.get('is_new'):
                new_player_count += 1
        
        return result
    
    def _diversify_by_game(self, candidates: List[Dict]) -> List[Dict]:
        game_groups = defaultdict(list)
        for c in candidates:
            game = c.get('games', ['其他'])[0]
            game_groups[game].append(c)
        
        result = []
        max_len = max(len(g) for g in game_groups.values()) if game_groups else 0
        
        for i in range(max_len):
            for game, players in game_groups.items():
                if i < len(players):
                    result.append(players[i])
        
        return result
    
    def _apply_operation_strategy(self, candidates: List[Dict], user_profile: Dict) -> List[Dict]:
        if len(candidates) < 10:
            return candidates
        
        result = candidates.copy()
        
        new_players = [p for p in self.players.values() if p.get('is_new')]
        if new_players and len(result) >= 5:
            new_player = new_players[0].copy()
            new_player['reason'] = '新人推荐'
            result.insert(4, new_player)
        
        high_quality = [
            p for p in self.players.values()
            if p.get('rating', 0) >= 4.8 and p.get('price', 0) >= 45
        ]
        if high_quality and len(result) >= 10:
            hq_player = high_quality[0].copy()
            hq_player['reason'] = '优质推荐'
            result.insert(9, hq_player)
        
        return result
    
    def recommend_by_game(
        self,
        user_id: int,
        game_id: int,
        filters: Optional[Dict] = None,
        page_size: int = 20
    ) -> List[Dict]:
        game_map = {1: '王者荣耀', 2: 'LOL', 3: '和平精英'}
        game_name = game_map.get(game_id, '王者荣耀')
        
        candidates = self._get_players_by_game(game_name, top_k=100)
        
        if filters:
            if 'price_min' in filters and 'price_max' in filters:
                candidates = [
                    p for p in candidates
                    if filters['price_min'] <= p.get('price', 0) <= filters['price_max']
                ]
            
            if 'voice_type' in filters:
                candidates = [
                    p for p in candidates
                    if p.get('voice_type') == filters['voice_type']
                ]
        
        candidates = self._apply_filters(candidates, user_id)
        
        user_profile = self._get_user_profile(user_id)
        ranked = self._rank_candidates(candidates, user_profile)
        
        return ranked[:page_size]
    
    def recommend_similar(self, player_id: int, count: int = 10) -> List[Dict]:
        if player_id not in self.player_embeddings:
            return []
        
        target_player = self.players[player_id]
        target_vec = self.player_embeddings[player_id]
        
        scores = {}
        for pid, pvec in self.player_embeddings.items():
            if pid == player_id:
                continue
            
            score = np.dot(target_vec, pvec) / (
                np.linalg.norm(target_vec) * np.linalg.norm(pvec)
            )
            
            other_player = self.players[pid]
            
            game_sim = len(set(target_player.get('games', [])) & set(other_player.get('games', [])))
            price_diff = abs(target_player.get('price', 0) - other_player.get('price', 0))
            price_sim = 1.0 - min(price_diff / 50, 1.0)
            voice_sim = 1.0 if target_player.get('voice_type') == other_player.get('voice_type') else 0.5
            
            final_score = score * 0.6 + game_sim * 0.2 + price_sim * 0.1 + voice_sim * 0.1
            scores[pid] = final_score
        
        sorted_ids = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:count]
        
        results = []
        for pid, score in sorted_ids:
            player = self.players[pid].copy()
            player['similarity_score'] = float(score)
            results.append(player)
        
        return results
    
    def load_users(self, users: List[Dict]):
        for user in users:
            self.users[user['user_id']] = user
    
    def load_orders(self, orders: List[Dict]):
        self.orders = orders
        
        for order in orders:
            player_id = order['player_id']
            self.stats['player_orders'][player_id] += 1
            self.stats['hot_24h'][player_id] += 1


